#include "cuda_tracer.h"
#include "cuda_tracer.cuh"

#include "math/MathLib.h"
#include "scene/scene.h"
#include "scene/objects.h"
#include "utility/image.h"
#include "scene/texture.h"

using rtMath::vec3;
using rtMath::vec4;
using rtMath::Ray;

inline float3 Vec3ToFloat3(const vec3& v)
{
	return make_float3(v.x, v.y, v.z);
}

CudaTracer::CudaTracer()
	: m_MaximumTraceDepth(s_MaxTraceDepth)
	, m_DataPrepared(false)
	, m_DeviceSpheresPtr(0)
	, m_SpheresCount(0)
	, m_DeviceBoxesPtr(0)
	, m_BoxesCount(0)
	, m_DevicePlanesPtr(0)
	, m_PlanesCount(0)
	, m_DeviceMaterialsPtr(0)
	, m_DeviceLightsPtr(0)
	, m_LightsCount(0)
	, m_DeviceCameraPtr(0)
	, m_DeviceImagePtr(0)
	, m_ImageWidth(0)
	, m_ImageHeight(0)
	, m_TexturesCount(0)
{
}
CudaTracer::~CudaTracer()
{
	if (m_DeviceSpheresPtr)
	{
		cudaSafeCall(cudaFree(m_DeviceSpheresPtr));
		m_DeviceSpheresPtr = 0;
	}
	if (m_DeviceBoxesPtr)
	{
		cudaSafeCall(cudaFree(m_DeviceBoxesPtr));
		m_DeviceBoxesPtr = 0;
	}
	if (m_DevicePlanesPtr)
	{
		cudaSafeCall(cudaFree(m_DevicePlanesPtr));
		m_DevicePlanesPtr = 0;
	}
	if (m_DeviceMaterialsPtr)
	{
		cudaSafeCall(cudaFree(m_DeviceMaterialsPtr));
		m_DeviceMaterialsPtr = 0;
	}
	if (m_DeviceLightsPtr)
	{
		cudaSafeCall(cudaFree(m_DeviceLightsPtr));
		m_DeviceLightsPtr = 0;
	}
	if (m_DeviceCameraPtr)
	{
		cudaSafeCall(cudaFree(m_DeviceCameraPtr));
		m_DeviceCameraPtr = 0;
	}
	if (m_DeviceImagePtr)
	{
		cudaSafeCall(cudaFree(m_DeviceImagePtr));
		m_DeviceImagePtr = 0;
	}

	for (int i = 0; i < m_TexturesCount; ++i)
		CudaFreeTexture(i);
}

void CudaTracer::SetMaximumTraceDepth(int maxTraceDepth)
{
	m_MaximumTraceDepth = maxTraceDepth;
}

bool CudaTracer::PrepareSceneData(const Scene& scene, Image& screen)
{
	int result = this->GrabSpheresFromScene(scene);
	result |= this->GrabBoxesFromScene(scene);
	result |= this->GrabPlanesFromScene(scene);
	result |= this->GrabLightsFromScene(scene);
	result |= this->GrabMaterialsFromScene(scene);

	if (!result)
		return false;

	this->GrabTexturesFromScene(scene);

	m_ImageWidth = screen.GetWidth();
	m_ImageHeight = screen.GetHeight();
	const int imgDataSize = m_ImageWidth * m_ImageHeight * sizeof(int);
	cudaSafeCall(cudaMalloc(&m_DeviceImagePtr, imgDataSize));

	SceneCamera* camera = scene.GetCamera();
	camera->Setup(static_cast<float>(m_ImageWidth), static_cast<float>(m_ImageHeight));
	this->UpdateCamera(camera);

	return true;
}

void CudaTracer::Render(Image& screen) const
{
	CudaRayTrace(reinterpret_cast<uint32*>(m_DeviceImagePtr), m_ImageWidth, m_ImageHeight,
				 reinterpret_cast<CudaSphere*>(m_DeviceSpheresPtr), m_SpheresCount,
				 reinterpret_cast<CudaBox*>(m_DeviceBoxesPtr), m_BoxesCount,
				 reinterpret_cast<CudaPlane*>(m_DevicePlanesPtr), m_PlanesCount,
				 reinterpret_cast<CudaMaterial*>(m_DeviceMaterialsPtr),
				 reinterpret_cast<CudaCamera*>(m_DeviceCameraPtr),
				 reinterpret_cast<CudaLight*>(m_DeviceLightsPtr), m_LightsCount);

	const int imgDataSize = screen.GetWidth() * screen.GetHeight() * sizeof(int);
	cudaSafeCall(cudaMemcpy(screen.GetDataPtr(), m_DeviceImagePtr, imgDataSize, cudaMemcpyDeviceToHost));
}

void CudaTracer::ResizeScreen(int width, int height)
{
	if (width != m_ImageWidth || m_ImageHeight != height)
	{
		cudaSafeCall(cudaFree(m_DeviceImagePtr));

		m_ImageWidth = width;
		m_ImageHeight = height;
		const int imgDataSize = m_ImageWidth * m_ImageHeight * sizeof(int);
		cudaSafeCall(cudaMalloc(&m_DeviceImagePtr, imgDataSize));
	}
}

void CudaTracer::UpdateCamera(SceneCamera* camera)
{
	if (!m_DeviceCameraPtr)
		cudaSafeCall(cudaMalloc(&m_DeviceCameraPtr, sizeof(CudaCamera)));

	CudaCamera cudaCameraTmp;
	cudaCameraTmp.position = Vec3ToFloat3(camera->GetPosition());
	cudaCameraTmp.direction = Vec3ToFloat3(camera->GetDirection());
	cudaCameraTmp.U = Vec3ToFloat3(camera->GetU());
	cudaCameraTmp.V = Vec3ToFloat3(camera->GetV());
	cudaSafeCall(cudaMemcpy(m_DeviceCameraPtr, &cudaCameraTmp, sizeof(CudaCamera), cudaMemcpyHostToDevice));
}



int CudaTracer::GrabSpheresFromScene(const Scene& scene)
{
	std::vector<CudaSphere>	sphVec;
	for (int i = 0; i < scene.GetShapesCount(); ++i)
	{
		const SceneShape* shape = scene.GetShape(i);
		if (shape->GetShapeType() == SceneShape::ShapeSphere)
		{
			CudaSphere sph;
			sph.center = Vec3ToFloat3(shape->GetPosition());
			sph.radius = static_cast<const ShapeSphere*>(shape)->GetRadius();
			sph.matID  = shape->GetMaterialID();

			sphVec.push_back(sph);
		}
	}

	if ((m_SpheresCount = static_cast<int>(sphVec.size())) > 0)
	{
		cudaSafeCall(cudaMalloc(&m_DeviceSpheresPtr, m_SpheresCount * sizeof(CudaSphere)));
		cudaSafeCall(cudaMemcpy(m_DeviceSpheresPtr, &sphVec.front(), m_SpheresCount * sizeof(CudaSphere), cudaMemcpyHostToDevice));
	}

	return m_SpheresCount;
}

int CudaTracer::GrabBoxesFromScene(const Scene& scene)
{
	std::vector<CudaBox>	boxVec;
	for (int i = 0; i < scene.GetShapesCount(); ++i)
	{
		const SceneShape* shape = scene.GetShape(i);
		if (shape->GetShapeType() == SceneShape::ShapeBox)
		{
			CudaBox box;
			box.position = Vec3ToFloat3(shape->GetPosition());
			box.Minimum = Vec3ToFloat3(static_cast<const ShapeBox*>(shape)->GetMinimum());
			box.Maximum = Vec3ToFloat3(static_cast<const ShapeBox*>(shape)->GetMaximum());
			box.matID  = shape->GetMaterialID();

			boxVec.push_back(box);
		}
	}

	if ((m_BoxesCount = static_cast<int>(boxVec.size())) > 0)
	{
		cudaSafeCall(cudaMalloc(&m_DeviceBoxesPtr, m_BoxesCount * sizeof(CudaBox)));
		cudaSafeCall(cudaMemcpy(m_DeviceBoxesPtr, &boxVec.front(), m_BoxesCount * sizeof(CudaBox), cudaMemcpyHostToDevice));
	}

	return m_BoxesCount;
}

int CudaTracer::GrabPlanesFromScene(const Scene& scene)
{
	std::vector<CudaPlane>	planeVec;
	for (int i = 0; i < scene.GetShapesCount(); ++i)
	{
		const SceneShape* shape = scene.GetShape(i);
		if (shape->GetShapeType() == SceneShape::ShapePlane)
		{
			CudaPlane plane;
			plane.normal = Vec3ToFloat3(static_cast<const ShapePlane*>(shape)->GetNormal());
			plane.distance = static_cast<const ShapePlane*>(shape)->GetDistance();
			plane.uAxis = Vec3ToFloat3(static_cast<const ShapePlane*>(shape)->GetUAxis());
			plane.vAxis = Vec3ToFloat3(static_cast<const ShapePlane*>(shape)->GetVAxis());
			plane.matID  = shape->GetMaterialID();

			planeVec.push_back(plane);
		}
	}

	if ((m_PlanesCount = static_cast<int>(planeVec.size())) > 0)
	{
		cudaSafeCall(cudaMalloc(&m_DevicePlanesPtr, m_PlanesCount * sizeof(CudaPlane)));
		cudaSafeCall(cudaMemcpy(m_DevicePlanesPtr, &planeVec.front(), m_PlanesCount * sizeof(CudaPlane), cudaMemcpyHostToDevice));
	}

	return m_PlanesCount;
}

int CudaTracer::GrabMaterialsFromScene(const Scene& scene)
{
	std::vector<CudaMaterial>	matVec;
	for (int i = 0; i < scene.GetMaterialsCount(); ++i)
	{
		const Material* mat = scene.GetMaterial(i);
		
		CudaMaterial material;
		material.diffuse = Vec3ToFloat3(mat->GetDiffuse());
		material.specular = Vec3ToFloat3(mat->GetSpecular());
		material.specularLevel = mat->GetSpecularLevel();
		material.glossiness = mat->GetGlossiness();
		material.reflectivity = mat->GetReflectivity();
		material.transparency = mat->GetTransparency();
		material.refractionIndex = mat->GetRefractionIndex();
		material.albedo = mat->GetAlbedoTexture();
		material.bump = mat->GetBumpTexture();

		matVec.push_back(material);
	}

	const int numMaterials = static_cast<int>(matVec.size());
	if (numMaterials > 0)
	{
		cudaSafeCall(cudaMalloc(&m_DeviceMaterialsPtr, sizeof(CudaMaterial) * numMaterials));
		cudaSafeCall(cudaMemcpy(m_DeviceMaterialsPtr, &matVec.front(), numMaterials * sizeof(CudaMaterial), cudaMemcpyHostToDevice));
	}

	return numMaterials;
}

int CudaTracer::GrabTexturesFromScene(const Scene& scene)
{
	const int numTextures = scene.GetTexturesCount();

	for (int i = 0; i < numTextures; ++i)
	{
		const Texture* texture = scene.GetTexture(i);
		CudaInitTexture(texture->GetDataPtr(), texture->GetWidth(), texture->GetHeight(), i);
	}

	return numTextures;
}

int CudaTracer::GrabLightsFromScene(const Scene& scene)
{
	std::vector<CudaLight>	lightVec;
	for (int i = 0; i < scene.GetLightsCount(); ++i)
	{
		const SceneLight* lght = scene.GetLight(i);

		CudaLight light;
		light.position = Vec3ToFloat3(lght->GetPosition());
		light.intensity = Vec3ToFloat3(lght->GetIntensity());

		lightVec.push_back(light);
	}

	if ((m_LightsCount = static_cast<int>(lightVec.size())) > 0)
	{
		cudaSafeCall(cudaMalloc(&m_DeviceLightsPtr, sizeof(CudaLight) * m_LightsCount));
		cudaSafeCall(cudaMemcpy(m_DeviceLightsPtr, &lightVec.front(), m_LightsCount * sizeof(CudaLight), cudaMemcpyHostToDevice));
	}

	return m_LightsCount;
}
